#include "SensorVisualisationHandler.h"

#include <VirtualRobot/RobotConfig.h>
#include <MMMSimoxTools/MMMSimoxTools.h>

SensorVisualisationHandler::SensorVisualisationHandler(SoSeparator* sceneGraph, bool interpolationEnabled, float timestepDelta)
    : sensorVisualisationSeperator(new SoSeparator()),
      interpolationEnabled(interpolationEnabled),
      timestepDelta(timestepDelta),
      sensorVisualisationTable(new SensorVisualisationTable())
{
    sceneGraph->addChild(sensorVisualisationSeperator);
    sensorVisualisationSeperator->ref();
}

void SensorVisualisationHandler::updateSensorVisualisation(const std::map<std::string, boost::shared_ptr<MMM::SensorVisualisationFactory> > &sensorVisualisationFactories) {
    this->sensorVisualisationFactories = sensorVisualisationFactories;
    loadSensorVisualisation(motions);
}

void SensorVisualisationHandler::loadSensorVisualisation(MMM::MotionList motions) {
    if (motions.size() == 0) return;
    bool visualisationLoaded = false;

    this->motions = motions;

    SoSeparator* sep = new SoSeparator();
    sep->ref();

    sensorVisualisations.clear();

    for (MMM::MotionPtr motion : motions)
    {
        if (motion->getSensorData().size() == 0) break;

        VirtualRobot::RobotPtr robot;
        VirtualRobot::CoinVisualizationPtr visualization;
        if (motion->getModel()) {
            robot = MMM::SimoxTools::buildModel(motion->getModel());
            MMM::SimoxTools::updateInertialMatricesFromModels(robot);
            visualization = robot->getVisualization<VirtualRobot::CoinVisualization>(VirtualRobot::SceneObject::Full);
        }

        std::map<std::string, MMM::SensorVisualisationPtr> motionSensorVis;
        for (auto sensor : motion->getPrioritySortedSensorData()) {
            MMM::SensorVisualisationFactoryPtr factory = sensorVisualisationFactories[MMM::VIS_STR(sensor->getType().c_str())];
            if (factory) {
                MMM::SensorVisualisationPtr visualisation = factory->createSensorVisualisation(sensor, robot, visualization, sep);
                if (visualisation) motionSensorVis[sensor->getUniqueName()] = visualisation;
            } else {
                MMM_INFO << "No sensor visualisation plugin for sensor " << sensor->getType() << " of motion " << motion->getName() << std::endl;
            }
        }

        if (motionSensorVis.size() > 0) {
            sensorVisualisationSeperator->removeAllChildren();
            sensorVisualisationSeperator->addChild(sep);
            sensorVisualisations[motion->getName()] = motionSensorVis;
            visualisationLoaded = true;
        } else {
            MMM_INFO << "No visualisation for motion " << motion->getName() << " possible!" << std::endl;
        }
    }

    if (visualisationLoaded) {
        sensorVisualisationTable->populateTable(motions, sensorVisualisations);
    } else {
        std::string message = "No motion visualisation!";
        sensorVisualisationTable->displayMessage(message);
        MMM_ERROR << message << " Check needed sensor and sensor visualisation plugins!" << std::endl;
    }

    emit sensorVisualisationLoaded(visualisationLoaded);
}

bool SensorVisualisationHandler::changeTimestep(float timestep) {
    if (motions.size() == 0) return false;

    for (MMM::MotionPtr motion : motions)
    {
        // TODO: Sort only once, not every timestep!
        std::vector<MMM::SensorPtr> sortedSensors = motion->getPrioritySortedSensorData();
        // Set the Kinematicsensors first (needed for some visualisations)!
        std::stable_sort(sortedSensors.begin(), sortedSensors.end(), [] (MMM::SensorPtr p1, MMM::SensorPtr p2) {
            if (p1->getType() == "Kinematic" && p2->getType() != "Kinematic") return 1;
            else if (p1->getType() != "Kinematic" && p2->getType() == "Kinematic") return -1;
            else return 0;
        });

        for (auto sensor : sortedSensors) {
            MMM::SensorVisualisationPtr visualisation = sensorVisualisations[motion->getName()][sensor->getUniqueName()];
            if (visualisation) {
                if (interpolationEnabled) visualisation->updateVisualisation(timestep);
                else visualisation->updateVisualisation(timestep, timestepDelta);
            }
        }
    }

    return true;
}

void SensorVisualisationHandler::setInterpolation(bool enabled) {
    interpolationEnabled = enabled;
}

void SensorVisualisationHandler::setTimestepDelta(float delta) {
    timestepDelta = delta;
}

QTreeWidget* SensorVisualisationHandler::getSensorVisualisationTable() {
    return sensorVisualisationTable;
}

void SensorVisualisationHandler::displaySensorVisualisationTable(bool display) {
    if (display) sensorVisualisationTable->show();
    else sensorVisualisationTable->hide();
}
